package com.hx.voice.recognition.gateway.handler;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.handler.AsyncPredicate;
import org.springframework.cloud.gateway.handler.predicate.AbstractRoutePredicateFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.function.Predicate;

import static org.springframework.cloud.gateway.filter.AdaptCachedBodyGlobalFilter.CACHED_REQUEST_BODY_KEY;

/**
 * @descripe：缓存请求的消息头与消息体
 * @company：南京华讯方舟通讯设备有限公司
 * @author：zhoujinbing
 * @time：2019/6/11 16:15
 * @version：V1.0
 */
@Slf4j
public class RequestBodyRoutePredicateFactory extends AbstractRoutePredicateFactory<RequestBodyRoutePredicateFactory.Config> {

    private static final List<HttpMessageReader<?>> messageReader = HandlerStrategies.withDefaults().messageReaders();
    public static final String REQUEST_BODY_ATTR = "requestBodyAttr";


    public RequestBodyRoutePredicateFactory() {
        super(Config.class);
    }

    @Override
    public AsyncPredicate<ServerWebExchange> applyAsync(Config config) {
        return exchange -> {
            if(!"POST".equals(exchange.getRequest().getMethodValue()))
            {
                return Mono.just(true);
            }

            Object cacheBody = exchange.getAttribute(REQUEST_BODY_ATTR);
            if(cacheBody != null)
            {
                try {
                    return Mono.just(true);
                }catch (ClassCastException e)
                {
                    if (log.isDebugEnabled()) {
                        log.debug("Predicate test failed because class in predicate does not match the cached body object",
                                e);
                    }
                }
                return Mono.just(true);
            }
            else
            {
                return DataBufferUtils.join(exchange.getRequest().getBody())
                        .flatMap(dataBuffer -> {
                            DataBufferUtils.retain(dataBuffer);
                            Flux<DataBuffer> cachedFlux = Flux.defer(() -> Flux.just(dataBuffer.slice(0, dataBuffer.readableByteCount())));
                            ServerHttpRequestDecorator mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                                @Override
                                public Flux<DataBuffer> getBody()
                                {
                                    return cachedFlux;
                                }
                            };
                            return ServerRequest.create(exchange.mutate().request(mutatedRequest).build(),messageReader)
                                    .bodyToMono(String.class)
                                    .doOnNext(objectValue ->{
                                        log.info("异步存放请求消息体数据:{}success",objectValue);
                                        exchange.getAttributes().put(REQUEST_BODY_ATTR,objectValue);
                                        exchange.getAttributes().put(CACHED_REQUEST_BODY_KEY,cachedFlux);
                                    } )
                                    .map(objectValue -> true);

                        });
            }
        };
    }

    @Override
    public Predicate<ServerWebExchange> apply(Config config) {
        throw new UnsupportedOperationException("ReadBodyPredicateFactory is only async");
    }

    public static class Config {

        private String attr;

        public String getAttr()
        {
            return this.attr;
        }

        public void setAttr(String attr)
        {
            this.attr = attr;
        }
    }
}